import io
import cv2
import insightface
from insightface import model_zoo
from insightface.app.face_analysis import FaceAnalysis
import numpy as np
from PIL import Image

assert insightface.__version__ >= '0.7'


class InsightFace():
    def __init__(self, models_path=None):
        self.model = FaceAnalysis(models_path=models_path)
        # gpu_id: 正数为GPU的ID，负数为使用CPU
        self.model.prepare(ctx_id=0, det_size=(640, 640))
        # 换脸的模型
        self.swapper = model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)

    # 加载视频流的一帧图片
    def load_and_process_frame(self, frame):
        # 将当前帧转换为InsightFace接受的格式并进行人脸检测
        img = frame  # 假设模型可以直接处理cv2.imread返回的BGR格式
        faces = self.model.get(img)
        return img, faces

    # 加载本地图片
    def load_and_process_image(self, image_path):
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像：{image_path}")
        faces = self.model.get(img)
        return img, faces

    # 加载接口传输的图片流
    def load_and_process_image_from_stream(self, image_data):
        # 将二进制数据转换为BytesIO对象
        image_buffer = io.BytesIO(image_data)

        # 使用PIL库加载BytesIO对象，然后转换为OpenCV格式
        pil_image = Image.open(image_buffer)
        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)  # OpenCV默认使用BGR通道顺序

        if img is None or img.shape[0] == 0 or img.shape[1] == 0:
            raise ValueError(f"无法读取图像：{img}")

        # 使用模型处理图片
        faces = self.model.get(img)
        return img, faces

    # 给图片画上人脸框
    def draw(self, img, faces):
        return self.model.draw_on(img, faces)

    def crop_faces(self, img, faces):
        cropped_faces = []

        for i in range(len(faces)):
            face = faces[i]
            box = face.bbox.astype(int)

            # 获取人脸框左上角和右下角坐标
            x1, y1 = box[0], box[1]
            x2, y2 = box[2], box[3]

            # 裁剪出人脸区域
            face_img = img[y1:y2, x1:x2]

            # 将裁剪出来的人脸区域添加到结果列表中
            cropped_faces.append(face_img)

        return cropped_faces

    # 获取提取的人脸特征
    def extract_face_features(self, faces):
        feats = [face.normed_embedding for face in faces]
        return np.array(feats, dtype=np.float32)

    # 计算欧氏距离
    def compute_nearest_neighbor(self, face1_feats, face2_feats):
        euclidean_distances = np.linalg.norm(face1_feats - face2_feats, axis=1)
        # print(euclidean_distances)
        return np.argmin(euclidean_distances), np.min(euclidean_distances)

    # 换脸
    def swap_face(self, img_body, faces_body, faces_source, face_body_nth=-1, face_source_nth=0):
        if face_body_nth > len(faces_body):
            raise IndexError(f"Invalid face_body_nth index. Expected value exceeds {len(faces_body)}")
        if face_source_nth > len(faces_source):
            raise IndexError(f"Invalid face_source_nth index. Expected value exceeds {len(faces_source)}")

        # 对人脸顺序排序（从左到右）
        faces_body = sorted(faces_body, key=lambda x: x.bbox[0])
        faces_source = sorted(faces_source, key=lambda x: x.bbox[0])
        # 选取第n个人脸
        faces_source = faces_source[face_source_nth]
        res = img_body.copy()

        # -1全部替换
        if face_body_nth < 0:
            # 所有人脸都替换
            for face in faces_body:
                res = self.swapper.get(res, face, faces_source, paste_back=True)
        else:
            face = faces_body[face_body_nth]
            res = self.swapper.get(res, face, faces_source, paste_back=True)

        return res
